# 构建神经网络
from __future__ import print_function
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim


class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        # 定义第一层卷积,输入维度/深度=1，输出维度=6，卷积核/过滤器 3*3
        # 输出结果的宽*高为（原始宽度-过滤器宽度+1）*（原始高度-过滤器高度+1）
        # Conv2d需要4D tensor
        self.conv1 = nn.Conv2d(1, 6, 3)
        # 定义第二层卷积,输入维度=6，输出维度=16，卷积核3*3
        self.conv2 = nn.Conv2d(6, 16, 3)
        # 定义全连接层
        self.fc1 = nn.Linear(16 * 6 * 6, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        # 任意卷积层后需要加上 激活层relu 和 池化层max_pool2d
        # 2*2的max池化层选取2*2的格子里最大的那个值
        x = F.max_pool2d(F.relu(self.conv1(x)), (2, 2))
        x = F.max_pool2d(F.relu(self.conv2(x)), 2)

        # print(x.size())
        # torch.Size([1, 16, 6, 6])

        # 卷积处理后，需要调整张量的形状
        x = x.view(-1, self.num_flat_features(x))

        # 全连接
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

    def num_flat_features(self, x):
        size = x.size()[1:]
        num_features = 1
        for s in size:
            num_features *= s
        return num_features


net = Net()
# print(net)

# net里的所有可训练参数
params = list(net.parameters())
# print(len(params))
# print(params[0].size())

input = torch.randn(1, 1, 32, 32)
out = net(input)
print(out)
print(out.size())

# 损失函数
target = torch.randn(10)
# 改变target的形状为二维张量
target = target.view(1, -1)
# 计算均方误差
criterion = nn.MSELoss()
loss = criterion(out, target)
print(loss)

# 反向传播
# 梯度清零
net.zero_grad()
print(net.conv1.bias.grad)
# loss.backward()
print(net.conv1.bias.grad)

# 更新网络参数
# 传入可训练参数和学习率
optimizer = optim.SGD(net.parameters(), lr=0.01)
optimizer.zero_grad()
loss.backward()
optimizer.step()
